import os
from transformers import AutoTokenizer, AutoModelForMaskedLM


model_path = r'C:\Users\admin\.cache\huggingface\hub\models--bert-base-chinese\snapshots\8d2a91f91cc38c96bb8b4556ba70c392f8d5ee55'


assert os.path.exists(model_path), '获取预训练模型失败！`model_path`不存在？'


def get_tokenizer():
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    # pretrained_model = AutoModelForMaskedLM.from_pretrained(model_path)
    return tokenizer


def get_pretrained_model__bert():
    from transformers import BertModel
    pretrained_model = BertModel.from_pretrained(model_path)
    # pretrained_model = AutoModelForMaskedLM.from_pretrained(model_path)
    return pretrained_model


if __name__ == '__main__':
    tokenizer = get_tokenizer()
    pretrained_model = get_pretrained_model__bert()









